import cv2
import numpy as np
import torch

__VERSION__ = "v1.0.0"

class ImageContourGeneratorNode:
    CATEGORY = "AmorForge"
    RETURN_TYPES = ("IMAGE",)
    RETURN_NAMES = ("image_with_contour",)
    OUTPUT_NODE = False
    FUNCTION = "generate_contour"

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image": ("IMAGE",),
                "mask": ("MASK",),
                "contour_color": (["red", "green", "blue", "white", "black"], {
                    "default": "red"
                }),
                "contour_thickness": ("INT", {
                    "default": 15,
                    "min": 1,
                    "max": 30,
                    "step": 1,
                    "display": "number"
                }),
                "blur_radius": ("INT", {
                    "default": 2,
                    "min": 1,
                    "max": 20,
                    "step": 2,
                    "display": "number"
                }),
            },
        }

    def generate_contour(self, image, mask, contour_color="red", contour_thickness=3, blur_radius=5):
        """Generate image with red contour based on mask for Gemini model reference"""
        try:
            # Convert tensor to numpy arrays
            image_np = image.cpu().numpy()
            mask_np = mask.cpu().numpy()

            # Handle different input shapes
            if len(image_np.shape) == 4:
                image_np = image_np[0]  # Get first image from batch

            if len(mask_np.shape) == 3:
                mask_np = mask_np[0]  # Get first mask from batch

            # Convert to uint8 format (0-255)
            image_uint8 = (image_np * 255).astype(np.uint8)

            # Convert from RGB to BGR for OpenCV processing
            image_bgr = cv2.cvtColor(image_uint8, cv2.COLOR_RGB2BGR)

            # Ensure mask is binary (0 or 255)
            mask_binary = (mask_np * 255).astype(np.uint8)

            # Apply Gaussian blur to mask if blur_radius > 1
            if blur_radius > 1:
                mask_binary = cv2.GaussianBlur(mask_binary, (blur_radius, blur_radius), 0)

            # Find contours
            contours, _ = cv2.findContours(mask_binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

            # Define color mapping (BGR format for OpenCV)
            color_map = {
                "red": (0, 0, 255),    # Blue=0, Green=0, Red=255 -> Red
                "green": (0, 255, 0),  # Blue=0, Green=255, Red=0 -> Green
                "blue": (255, 0, 0),   # Blue=255, Green=0, Red=0 -> Blue
                "white": (255, 255, 255),
                "black": (0, 0, 0)
            }

            contour_color_bgr = color_map.get(contour_color, (0, 0, 255))  # Default to red

            # Draw contours on image
            image_with_contour = image_bgr.copy()
            cv2.drawContours(image_with_contour, contours, -1, contour_color_bgr, contour_thickness)

            # Convert back from BGR to RGB for ComfyUI
            image_rgb = cv2.cvtColor(image_with_contour, cv2.COLOR_BGR2RGB)

            # Convert back to float32 format (0-1) for ComfyUI
            image_with_contour_float = image_rgb.astype(np.float32) / 255.0

            # Convert to tensor (B,H,W,C) format
            image_tensor = torch.from_numpy(image_with_contour_float)[None,]

            return (image_tensor,)

        except Exception as e:
            error_msg = f"Contour generation failed: {str(e)}"
            print(f"[ImageContourGeneratorNode] Error: {error_msg}")
            raise Exception(error_msg)